import math
import time
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
from scipy import ndimage
import pickle
import os

class MapGenerator3D:
    """
    高效的三维栅格与 ESDF 生成器。

    两种建图方式：
      - build_from_box(...) : 直接在栅格层面把外部边界（箱体）和若干内部面/墙设为占据（非常快）
      - build_from_point_cloud(ox, oy, oz) : 将稀疏/任意点云索引到栅格，然后做膨胀（适用于真实点云）

    最后使用 ndimage.distance_transform_edt 计算 ESDF，并用距离阈值执行 robot 半径膨胀。
    """

    def __init__(self, resolution: float, robot_radius: float, ref_lat=0.0, ref_lon=0.0):
        """
        resolution : 栅格分辨率（米）
        robot_radius : 机器人半径（米），用于膨胀障碍
        """
        self.resolution = float(resolution)
        self.robot_radius = float(robot_radius)
        # 地图的地理参考点
        self.ref_lat = ref_lat
        self.ref_lon = ref_lon
        # 世界坐标边界（在 build 时计算）
        self.min_x = self.min_y = self.min_z = 0.0
        self.max_x = self.max_y = self.max_z = 0.0

        # 栅格维度（在 build 时计算）
        self.x_width = self.y_width = self.z_width = 0

        # 数据（在 build 时填充）
        self.obs_map = None   # bool ndarray: True 表示占据
        self.esdf_map = None  # float ndarray: 每个栅格到最近障碍的欧式距离 (m)

    # ---------------- 辅助转换 ----------------
    def calc_grid_position(self, index: int, min_position: float) -> float:
        """栅格索引 -> 世界坐标（格子的中心）"""
        return index * self.resolution + min_position

    def calc_xyz_index(self, pos: float, min_pos: float) -> int:
        """世界坐标 -> 栅格索引（round）"""
        return int(round((pos - min_pos) / self.resolution))

    # ---------------- 建图方法 ----------------
    def _setup_bounds(self, min_x, max_x, min_y, max_y, min_z, max_z):
        """内部：计算边界与栅格尺寸"""
        self.min_x = float(math.floor(min_x))
        self.min_y = float(math.floor(min_y))
        self.min_z = float(math.floor(min_z))
        self.max_x = float(math.ceil(max_x))
        self.max_y = float(math.ceil(max_y))
        self.max_z = float(math.ceil(max_z))

        self.x_width = max(1, int(math.ceil((self.max_x - self.min_x) / self.resolution)))
        self.y_width = max(1, int(math.ceil((self.max_y - self.min_y) / self.resolution)))
        self.z_width = max(1, int(math.ceil((self.max_z - self.min_z) / self.resolution)))

        # 初始化占据格子数组
        self.obs_map = np.zeros((self.x_width, self.y_width, self.z_width), dtype=bool)
        self.esdf_map = None

    # ---------------- 新增：生成空地图 ----------------
    def build_empty_box(self, x_min, x_max, y_min, y_max, z_min, z_max, esdf_fill_value=1e6):
        """
        生成一个“完全无障碍”的地图：
          - obs_map 全为 False（没有占据）
          - esdf_map 全为 esdf_fill_value（非常大，远大于任何 SAFE_DIS）
        适合你“先生成空地图，之后加载用于规划”的场景。
        """
        self._setup_bounds(x_min, x_max, y_min, y_max, z_min, z_max)
        # obs_map 已由 _setup_bounds 初始化为全 False
        # 直接把 ESDF 填为大值，避免调用 distance_transform（对全空可能返回不靠谱值）
        self.esdf_map = np.full((self.x_width, self.y_width, self.z_width), float(esdf_fill_value), dtype=float)
        return self
    
    def build_from_box(self, x_min, x_max, y_min, y_max, z_min, z_max, add_walls=True, internal_walls=None):
        """
        快速根据一个轴对齐长方体（箱体）构建地图（非常适合你原先的箱体边界场景）。

        - add_walls=True: 将六个边界面（最外围的 voxel 层）设为占据（底/顶/四壁）
        - internal_walls: 可选列表，元素为 dict 描述内部障碍（例如斜墙），每个 dict 形如：
              {"type":"plane_xy","z":10, "x_start":0,"x_end":30,"y_start":0,"y_end":30}
          我同时支持简单的点列表形式： {"points": [(x,y,z), ...]}
        """
        # 设定边界并初始化栅格
        self._setup_bounds(x_min, x_max, y_min, y_max, z_min, z_max)

        if add_walls:
            # 最外层 voxel 作为占据（厚度为 1 voxel）
            self.obs_map[0, :, :] = True          # x = min face
            self.obs_map[-1, :, :] = True         # x = max face
            self.obs_map[:, 0, :] = True          # y = min face
            self.obs_map[:, -1, :] = True         # y = max face
            self.obs_map[:, :, 0] = True          # z = min face (底)
            self.obs_map[:, :, -1] = True         # z = max face (顶)

        # 处理内部简单障碍（如果有）
        if internal_walls:
            for wall in internal_walls:
                if 'points' in wall:
                    # 直接用点列表（稀疏）
                    pts = np.array(wall['points'])
                    if pts.size > 0:
                        xi = np.round((pts[:, 0] - self.min_x) / self.resolution).astype(int)
                        yi = np.round((pts[:, 1] - self.min_y) / self.resolution).astype(int)
                        zi = np.round((pts[:, 2] - self.min_z) / self.resolution).astype(int)
                        xi = np.clip(xi, 0, self.x_width - 1)
                        yi = np.clip(yi, 0, self.y_width - 1)
                        zi = np.clip(zi, 0, self.z_width - 1)
                        self.obs_map[xi, yi, zi] = True
                else:
                    # 支持 plane_xy/plane_xz/plane_yz 类型
                    t = wall.get('type', '')
                    if t == 'plane_xy':
                        z = wall.get('z', 0.0)
                        zi = self.calc_xyz_index(z, self.min_z)
                        xi0 = self.calc_xyz_index(wall.get('x_start', self.min_x), self.min_x)
                        xi1 = self.calc_xyz_index(wall.get('x_end', self.max_x), self.min_x)
                        yi0 = self.calc_xyz_index(wall.get('y_start', self.min_y), self.min_y)
                        yi1 = self.calc_xyz_index(wall.get('y_end', self.max_y), self.min_y)
                        xi0, xi1 = np.clip([xi0, xi1], 0, self.x_width - 1)
                        yi0, yi1 = np.clip([yi0, yi1], 0, self.y_width - 1)
                        zi = int(np.clip(zi, 0, self.z_width - 1))
                        self.obs_map[xi0:xi1+1, yi0:yi1+1, zi] = True
                    elif t == 'plane_xz':
                        y = wall.get('y', 0.0)
                        yi = self.calc_xyz_index(y, self.min_y)
                        xi0 = self.calc_xyz_index(wall.get('x_start', self.min_x), self.min_x)
                        xi1 = self.calc_xyz_index(wall.get('x_end', self.max_x), self.min_x)
                        zi0 = self.calc_xyz_index(wall.get('z_start', self.min_z), self.min_z)
                        zi1 = self.calc_xyz_index(wall.get('z_end', self.max_z), self.min_z)
                        xi0, xi1 = np.clip([xi0, xi1], 0, self.x_width - 1)
                        zi0, zi1 = np.clip([zi0, zi1], 0, self.z_width - 1)
                        yi = int(np.clip(yi, 0, self.y_width - 1))
                        self.obs_map[xi0:xi1+1, yi, zi0:zi1+1] = True
                    elif t == 'plane_yz':
                        x = wall.get('x', 0.0)
                        xi = self.calc_xyz_index(x, self.min_x)
                        yi0 = self.calc_xyz_index(wall.get('y_start', self.min_y), self.min_y)
                        yi1 = self.calc_xyz_index(wall.get('y_end', self.max_y), self.min_y)
                        zi0 = self.calc_xyz_index(wall.get('z_start', self.min_z), self.min_z)
                        zi1 = self.calc_xyz_index(wall.get('z_end', self.max_z), self.min_z)
                        yi0, yi1 = np.clip([yi0, yi1], 0, self.y_width - 1)
                        zi0, zi1 = np.clip([zi0, zi1], 0, self.z_width - 1)
                        xi = int(np.clip(xi, 0, self.x_width - 1))
                        self.obs_map[xi, yi0:yi1+1, zi0:zi1+1] = True

        # 膨胀并计算 ESDF（用距离场而非循环）
        self._finalize_obs_and_esdf()
        return self

    def build_from_point_cloud(self, ox, oy, oz):
        """
        将任意点云投影到栅格上（适用于稀疏点云）。
        ox, oy, oz: 可迭代的坐标列表或 numpy 数组（单位：米）
        """
        if len(ox) == 0:
            raise ValueError("点云为空")

        ox_a = np.asarray(ox, dtype=float)
        oy_a = np.asarray(oy, dtype=float)
        oz_a = np.asarray(oz, dtype=float)

        # 计算边界时稍微扩展一小段，避免边界问题
        padding = max(self.resolution, self.robot_radius) + 1.0
        min_x = ox_a.min() - padding
        max_x = ox_a.max() + padding
        min_y = oy_a.min() - padding
        max_y = oy_a.max() + padding
        min_z = oz_a.min() - padding
        max_z = oz_a.max() + padding

        self._setup_bounds(min_x, max_x, min_y, max_y, min_z, max_z)

        # 将点云索引化到栅格（向量化）
        xi = np.round((ox_a - self.min_x) / self.resolution).astype(int)
        yi = np.round((oy_a - self.min_y) / self.resolution).astype(int)
        zi = np.round((oz_a - self.min_z) / self.resolution).astype(int)

        xi = np.clip(xi, 0, self.x_width - 1)
        yi = np.clip(yi, 0, self.y_width - 1)
        zi = np.clip(zi, 0, self.z_width - 1)

        self.obs_map[xi, yi, zi] = True

        # 膨胀并计算 ESDF
        self._finalize_obs_and_esdf()
        return self

    def _finalize_obs_and_esdf(self):
        """
        给定已标记的初始占据 self.obs_map（True = 原障碍或边界），
        通过 distance transform 实现 robot_radius 膨胀，生成最终 obs_map 与 esdf_map。
        """
        if self.obs_map is None:
            raise RuntimeError("请先初始化 self.obs_map")

        # 计算自由空间的距离到最近占据体素（单位：格子数），乘以 resolution 得到米
        free_space = np.logical_not(self.obs_map)
        # distance_transform_edt 的返回值是以像素/格子为单位的距离
        dist_grid = ndimage.distance_transform_edt(free_space) * self.resolution

        # 将距离小于等于 robot_radius 的格子标为占据（膨胀）
        inflated = dist_grid <= self.robot_radius
        self.obs_map = np.array(inflated, dtype=bool)

        # 最终 ESDF（对膨胀后的 free space 计算距离）
        free_after = np.logical_not(self.obs_map)
        self.esdf_map = ndimage.distance_transform_edt(free_after) * self.resolution

    # ---------------- 保存与可视化 ----------------
    def save_map(self, filename="3d_map_data.pkl"):
        """保存地图数据到文件（pickle）"""
        if self.obs_map is None or self.esdf_map is None:
            raise ValueError("请先构建地图再保存")

        map_data = {
            'ref_lat': self.ref_lat,
            'ref_lon': self.ref_lon,
            'resolution': self.resolution,
            'robot_radius': self.robot_radius,
            'min_x': self.min_x,
            'min_y': self.min_y,
            'min_z': self.min_z,
            'max_x': self.max_x,
            'max_y': self.max_y,
            'max_z': self.max_z,
            'x_width': self.x_width,
            'y_width': self.y_width,
            'z_width': self.z_width,
            'obs_map': self.obs_map,
            'esdf_map': self.esdf_map
        }
        with open(filename, 'wb') as f:
            pickle.dump(map_data, f)
        print(f"[save_map] 地图已保存到: {filename}")
        return filename

    def visualize_map(self, save_path=None, show=True, subsample=1):
        """3D 可视化占据栅格（可选保存）。subsample 用于减少绘图点数（加速）"""
        if self.obs_map is None:
            raise ValueError("请先构建地图再可视化")

        xs, ys, zs = [], [], []
        # 采样输出避免绘图瓶颈
        for ix in range(0, self.x_width, subsample):
            for iy in range(0, self.y_width, subsample):
                for iz in range(0, self.z_width, subsample):
                    if self.obs_map[ix, iy, iz]:
                        xs.append(self.calc_grid_position(ix, self.min_x))
                        ys.append(self.calc_grid_position(iy, self.min_y))
                        zs.append(self.calc_grid_position(iz, self.min_z))

        # 确保创建新图形前，清除当前可能残留的绘图状态
        plt.clf()
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(111, projection='3d')
        if xs:
            ax.scatter(xs, ys, zs, c='k', s=6, alpha=0.6, label='obstacle voxels')
        ax.set_box_aspect((max(1, self.x_width), max(1, self.y_width), max(1, self.z_width)))
        ax.set_title('3D Occupancy Grid (voxels)')
        ax.set_xlabel('X [m]'); ax.set_ylabel('Y [m]'); ax.set_zlabel('Z [m]')
        ax.legend()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"[visualize_map] 可视化已保存到 {save_path}")
        if show:
            plt.show()
        return fig


# ---------------- 示例：快速生成粗/细地图 ----------------
if __name__ == "__main__":
    # 配置
    coarse_resolution = 10.0
    coarse_robot_radius = 8.0
    fine_resolution = 2.0
    fine_robot_radius = 1.0
    # 场景边界
    x_min, x_max = -2000, 2000
    y_min, y_max = -2000, 2000
    z_min, z_max = 0, 300

    # # 1) 快速生成“箱体”式障碍
    # t0 = time.time()
    # mg_coarse = MapGenerator3D(resolution=coarse_resolution, robot_radius=robot_radius)
    # # 内部可以传入一些示例斜墙
    # internal = [
    #     # # 例：水平面（xy）在 z=10，x 0..200, y 0..200
    #     # {"type": "plane_xy", "z": 10.0, "x_start": -100, "x_end": 100, "y_start": -50, "y_end": 50},
    #     # # 例：斜线用稀疏点描述（少量点，示范用）
    #     # {"points": [(0 + s, 5 + s//2, 10) for s in range(0, 40, 4)]}
    # ]
    # mg_coarse.build_from_box(x_min, x_max, y_min, y_max, z_min, z_max, add_walls=True, internal_walls=internal)
    # mg_coarse.save_map("3d_environment_coarse_map.pkl")
    # mg_coarse.visualize_map(save_path="3d_environment_coarse_visualization.png", subsample=1)
    # print(f"粗图生成耗时: {time.time() - t0:.2f}s\n")

    # # 2) 生成细网格（以相同边界为例）
    # t1 = time.time()
    # mg_fine = MapGenerator3D(resolution=fine_resolution, robot_radius=robot_radius)
    # mg_fine.build_from_box(x_min, x_max, y_min, y_max, z_min, z_max, add_walls=True, internal_walls=internal)
    # mg_fine.save_map("3d_environment_fine_map.pkl")
    # mg_fine.visualize_map(save_path="3d_environment_fine_visualization.png", subsample=1)
    # print(f"细图生成耗时: {time.time() - t1:.2f}s\n")

    # print("完成：粗/细地图已生成并保存。")

    ref_lat = 29.358  # 示例参考纬度
    ref_lon = 106  # 示例参考经度
    mgc = MapGenerator3D(resolution=coarse_resolution, robot_radius=coarse_robot_radius,ref_lat=ref_lat,ref_lon=ref_lon)
    mgc.build_empty_box(x_min, x_max, y_min, y_max, z_min, z_max)
    mgc.save_map("3d_environment_coarse_map.pkl")
    mgc.visualize_map("3d_environment_coarse_visualization.png")

    # mgf = MapGenerator3D(resolution=fine_resolution, robot_radius=fine_robot_radius)
    # mgf.build_empty_box(x_min, x_max, y_min, y_max, z_min, z_max)
    # mgf.save_map("C:\\Users\\machenike\\Desktop\\项目\\9院低空经济\\multi_uav_planner_3d_with_wayponits\\3d_environment_fine_map.pkl")
    # mgf.visualize_map("C:\\Users\\machenike\\Desktop\\项目\\9院低空经济\\multi_uav_planner_3d_with_wayponits\\3d_environment_fine_visualization.png")